"""
Moonshot API提供商
支持访问Moonshot的模型，包含完整的功能组件支持
"""
from typing import Dict, Any, List, Generator, Optional
import logging
import json
import time
import requests

from ..core.provider import RemoteProvider
from ..core.exceptions import map_provider_error
from ..implementations.moonshot.partial_mode import MoonshotPartialMode
from ..implementations.moonshot.tool_call import MoonshotToolCall
from ..implementations.moonshot.web_search import MoonshotWebSearch
from ..implementations.moonshot.streaming import MoonshotStreaming
from ..core.provider_registry import register_provider

# 设置日志
logger = logging.getLogger("llm.moonshot")

@register_provider("moonshot")
class MoonshotProvider(RemoteProvider):
	"""
	Moonshot API提供商

	支持功能:
	- 网络搜索
	- 工具调用
	- 流式输出
	- 部分输出续写
	"""

	# 可用模型常量
	MOONSHOT_V1_8K = 'moonshot-v1-8k'
	MOONSHOT_V1_32K = 'moonshot-v1-32k'
	MOONSHOT_V1_128K = 'moonshot-v1-128k'
	MOONSHOT_V1_AUTO = 'moonshot-v1-auto'
	KIMI_LATEST = 'kimi-latest'

	def __init__(self, config: Dict[str, Any]):
		"""
		初始化Moonshot提供商

		Args:
			config: 配置字典
		"""
		# 先设置关键属性，确保在_init_capabilities被调用前可用
		self.search = config.get("search", False)  # 是否启用搜索功能
		self._client = None

		# 请求超时设置(秒)
		self.timeout = config.get("timeout", 60)

		# 确保正确初始化所有父类
		super().__init__(config)

		# 初始化特有配置
		self.model = config.get("model", self.MOONSHOT_V1_AUTO)

		# 验证配置
		self._validate_config(config)

	def _validate_config(self, config: Dict[str, Any]) -> None:
		"""
		验证配置合法性

		Args:
			config: 配置字典

		Raises:
			ConfigurationError: 如果配置不合法
		"""
		# 验证模型名称
		model = config.get("model", self.MOONSHOT_V1_AUTO)
		valid_models = [
			self.MOONSHOT_V1_8K,
			self.MOONSHOT_V1_32K,
			self.MOONSHOT_V1_128K,
			self.MOONSHOT_V1_AUTO,
			self.KIMI_LATEST,
		]
		if model not in valid_models:
			logger.warning(f"未知的Moonshot模型: {model}，将使用默认模型 {self.MOONSHOT_V1_AUTO}")

	def _init_capabilities(self) -> None:
		"""初始化Moonshot支持的功能组件"""
		logger.debug("初始化Moonshot功能组件")

		# 添加部分输出功能
		self.register_capability("partial_mode", MoonshotPartialMode)

		# 添加工具调用功能
		self.register_capability("tool_calls", MoonshotToolCall)

		# 添加网络搜索功能（如果配置中启用）
		if self.search:
			self.register_capability("web_search", MoonshotWebSearch)

		# 添加流式输出功能
		self.register_capability("streaming", MoonshotStreaming)

	def filter_chat_args(self, **kwargs) -> tuple:
		"""
		过滤和处理API调用参数

		Args:
			**kwargs: 关键字参数

		Returns:
			处理后的(args, kwargs)元组
		"""
		# 添加搜索工具（保持向后兼容性）
		if self.search and "tools" not in kwargs:
			kwargs["tools"] = [{"type": "builtin_function", "function": {"name": "$web_search"}}]

		return (), kwargs

	def _get_client(self):
		"""
		获取Moonshot客户端

		尝试使用SDK客户端，如果不可用则返回None，回退到直接API调用

		Returns:
			Moonshot客户端实例或None
		"""
		if self._client is not None:
			return self._client
			
		try:
			from moonshot import Client

			# 创建客户端
			client = Client(api_key=self.api_key)

			# 设置HTTP客户端（如果有）
			http_client = self._setup_http_client()
			if http_client:
				client.http_client = http_client

			self._client = client
			logger.debug("创建了新的Moonshot客户端")
			return self._client
		except ImportError:
			logger.debug("Moonshot SDK未安装，将使用直接API调用")
			return None

	def _prepare_request_params(self, messages: List[Dict[str, Any]], **kwargs) -> Dict[str, Any]:
		"""
		准备API请求参数

		Args:
			messages: 消息列表
			**kwargs: 额外的提供商特定参数

		Returns:
			请求参数字典
		"""
		return {
			"model": self.model,
			"messages": messages,
			"temperature": self.temperature,
			"max_tokens": self.max_tokens,
			**kwargs
		}

	def _add_tools_to_request(self, params: Dict[str, Any]) -> Dict[str, Any]:
		"""
		添加工具功能到请求参数

		Args:
			params: 请求参数字典

		Returns:
			更新后的请求参数字典
		"""
		result = params.copy()

		# 添加工具功能（如果可用）
		if self.has_capability("tool_calls"):
			tool_call = self.get_capability("tool_calls")
			result["tools"] = tool_call.format_tools_for_request()

		# 添加网络搜索工具（如果可用）
		if self.has_capability("web_search"):
			web_search = self.get_capability("web_search")

			if "tools" in result:
				result["tools"].extend(web_search.prepare_search_tools())
			else:
				result["tools"] = web_search.prepare_search_tools()

		return result

	def _handle_tool_calls(self, completion: Any) -> List[Dict[str, Any]]:
		"""
		处理工具调用结果

		Args:
			completion: API完成响应

		Returns:
			工具执行结果列表
		"""
		results = []

		if not self.has_capability("tool_calls"):
			return results

		tool_call_capability = self.get_capability("tool_calls")
		tool_calls = tool_call_capability.detect_tool_calls(completion)

		if tool_calls:
			results = tool_call_capability.execute_tool_calls(tool_calls)

		return results

	def complete_stream(self, messages: List[Dict[str, str]]) -> Generator[str, None, None]:
		"""
		获取LLM流式完成
		
		Args:
			messages: 消息列表
			
		Returns:
			流式响应生成器
		"""
		# 检查是否支持流式输出
		if not self.has_capability("streaming"):
			# 如果不支持流式输出，使用标准完成方法并逐字符生成
			content = self._complete_internal(messages)
			for char in content:
				yield char
			return
			
		# 使用流式输出能力组件
		streaming = self.get_capability("streaming")
		yield from streaming.stream_response(messages)
		
	def get_current_messages(self) -> List[Dict[str, Any]]:
		"""
		获取当前会话消息
		用于部分输出模式
		
		Returns:
			消息列表，如果未保持状态则为空列表
		"""
		# 基础实现不保持状态，返回空列表
		# 子类可以重写此方法提供会话状态
		return []
		
	def _complete_internal(self, messages: List[Dict[str, Any]], **kwargs) -> str:
		"""
		内部完成方法实现

		Args:
			messages: 消息列表
			**kwargs: 额外参数

		Returns:
			完成的文本
		"""
		# 准备请求参数
		params = self._prepare_request_params(messages, **kwargs)
		
		# 添加工具（如果有）
		params = self._add_tools_to_request(params)
		
		# 过滤参数
		_, params = self.filter_chat_args(**params)

		# 首先尝试使用SDK客户端
		client = self._get_client()
		if client:
			try:
				# 使用客户端API调用
				response = self._call_with_retry(
					client.chat.completions.create,
					**params
				)

				# 获取文本响应
				choice = response.choices[0]
				content = choice.message.content

				# 处理工具调用（如果有）
				finish_reason = choice.finish_reason
				if finish_reason == "tool_calls":
					tool_results = self._handle_tool_calls(response)
					if tool_results:
						# 将工具结果添加到消息并重新调用API
						new_messages = messages.copy()
						new_messages.append({"role": "assistant", "content": content})
						for result in tool_results:
							new_messages.append(result)
						return self._complete_internal(new_messages)

				return content
			except Exception as e:
				logger.warning(f"SDK调用失败，回退到直接API调用: {str(e)}")
				# 回退到直接API调用
		
		# 直接API调用
		try:
			# 准备请求
			headers = {
				"Authorization": f"Bearer {self.api_key}",
				"Content-Type": "application/json"
			}

			# 使用适当的模型设置
			if "model" not in params:
				params["model"] = self.model

			# 确保流式设置为False
			params["stream"] = False

			# 发送请求
			url = "https://api.moonshot.cn/v1/chat/completions"
			response = self._call_with_retry(
				requests.post,
				url,
				headers=headers,
				json=params,
				timeout=self.timeout
			)

			# 检查响应状态
			if response.status_code != 200:
				raise Exception(f"API请求失败: 状态码 {response.status_code} 错误: {response.text}")

			# 解析响应
			data = response.json()
			choice = data["choices"][0]
			content = choice["message"]["content"]

			# 处理工具调用（如果有）
			finish_reason = choice.get("finish_reason")
			if finish_reason == "tool_calls" and "message" in choice and "tool_calls" in choice["message"]:
				tool_calls = choice["message"]["tool_calls"]
				
				# 保存assistant消息和工具调用
				new_messages = messages.copy()
				new_messages.append({
					"role": "assistant", 
					"content": content, 
					"tool_calls": tool_calls
				})
				
				# 处理每个工具调用
				for tool_call in tool_calls:
					tool_call_id = tool_call["id"]
					function = tool_call["function"]
					name = function["name"]
					arguments = json.loads(function["arguments"])
					
					# 为网络搜索创建工具响应
					if name == "$web_search" and self.has_capability("web_search"):
						web_search = self.get_capability("web_search")
						tool_message = {
							"role": "tool",
							"tool_call_id": tool_call_id,
							"name": "$web_search",
							"content": json.dumps(arguments, ensure_ascii=False)
						}
						new_messages.append(tool_message)
				
				# 再次请求，获取最终结果
				return self._complete_internal(new_messages)
				
			return content
		except Exception as e:
			# 映射错误为标准LLM异常
			mapped_error = map_provider_error("moonshot", e, self.model)
			raise mapped_error
